"""
Callback context and manager.
"""

import torch
import numpy as np
from typing import List, Optional

"""
Base callback class for Sponge.
"""

class Callback:
    """Base class for callbacks.
    
    Callbacks can be called at different stages during training:
    - begin: Called before the entire training starts
    - epoch_begin: Called before each epoch starts
    - step_begin: Called before each step starts
    - step_end: Called after each step ends
    - epoch_end: Called after each epoch ends
    - end: Called after the entire training ends
    
    Args:
        priority (int): Priority of the callback, smaller number means higher priority. Default: 1.
    """
    def __init__(self, priority=1):
        self._priority = priority

    @property
    def priority(self):
        """Get the priority of the callback"""
        return self._priority

    def begin(self, run_context):
        """Called before training starts.

        Args:
            run_context (RunContext): Context containing the callback running environment.
        """
        pass

    def epoch_begin(self, run_context):
        """Called before each epoch starts.

        Args:
            run_context (RunContext): Context containing the callback running environment.
        """
        pass

    def step_begin(self, run_context):
        """Called before each step starts.

        Args:
            run_context (RunContext): Context containing the callback running environment.
        """
        pass

    def step_end(self, run_context):
        """Called after each step ends.

        Args:
            run_context (RunContext): Context containing the callback running environment.
        """
        pass

    def epoch_end(self, run_context):
        """Called after each epoch ends.

        Args:
            run_context (RunContext): Context containing the callback running environment.
        """
        pass

    def end(self, run_context):
        """Called after training ends.

        Args:
            run_context (RunContext): Context containing the callback running environment.
        """
        pass 

class RunContext:
    """Runtime context for passing information and managing parameters in callbacks
    
    Contains:
    - Basic info: network, optimizer, dataset, etc.
    - Training progress: current step, epoch, etc.
    - Physical quantities: energy, force, etc.
    - Evaluation metrics, etc.
    """
    def __init__(self):
        # Basic info
        self.network = None  # Network
        self.optimizer = None  # Optimizer
        self.dataset = None  # Dataset
        self.parallel_mode = None  # Parallel mode
        self.device_number = None  # Number of devices
        
        # Training progress
        self.total_steps = 0  # Total number of steps
        self.cur_step = 0  # Current step
        self.cur_time = 0.0  # Current simulation time
        
        # Physical quantities
        self.potential = None  # Potential energy
        self.force = None  # Force
        self.energies = None  # Energy terms
        self.num_energies = 0  # Number of energy terms
        self.energy_names = []  # Names of energy terms
        self.energy_function = None  # Energy function
        self.bias = None  # Bias potential
        self.biases = None  # Bias potential terms
        self.bias_names = []  # Names of bias terms
        self.bias_function = None  # Bias function
        self.force_function = None  # Force function

        self.volume = None  # Volume
        self.coordinate = None  # Coordinates
        self.pbc_box = None  # Periodic boundary condition box
        self.image = None  # Image

        self.velocity = None  # Velocity
        self.kinetics = None  # Kinetic energy
        self.temperature = None  # Temperature
        self.pressure = None  # Pressure

        # Evaluation metrics
        self.metrics = None  # Metrics
        
        # Others
        self.list_callback = None  # List of callbacks
        self._stop_requested = False  # Whether stop is requested

    def __getattribute__(self, name):
        value = object.__getattribute__(self, name)
        if isinstance(value, torch.Tensor):
            return value.squeeze(-1).detach().cpu().numpy()
        return value

    def __setattr__(self, name, value):
        object.__setattr__(self, name, value)
    
    def update_basic_info(self, **kwargs):
        """Update basic information"""
        for k, v in kwargs.items():
            if hasattr(self, k):
                setattr(self, k, v)
                
    def update_progress(self, **kwargs):
        """Update training progress"""
        for k, v in kwargs.items():
            if hasattr(self, k):
                setattr(self, k, v)
                
    def update_physical_quantities(self, **kwargs):
        """Update physical quantities"""
        for k, v in kwargs.items():
            if hasattr(self, k):
                setattr(self, k, v)
                
    def update_metrics(self, **kwargs):
        """Update evaluation metrics"""
        for k, v in kwargs.items():
            if hasattr(self, k):
                setattr(self, k, v)

    def request_stop(self):
        """Request to stop running"""
        self._stop_requested = True

    def get_stop_requested(self):
        """Check if stop is requested"""
        return self._stop_requested

    # For backward compatibility
    def original_args(self):
        """Get original parameters (for backward compatibility)"""
        return self

class CallbackManager:
    """Callback manager
    
    Args:
        callbacks: List of callbacks
    """
    def __init__(self, callbacks=None):
        self._callbacks = []
        if callbacks:
            self._callbacks.extend(callbacks)

    def begin(self, run_context):
        """Execute begin callback"""
        for cb in self._callbacks:
            cb.begin(run_context)

    def epoch_begin(self, run_context):
        """Execute epoch_begin callback"""
        for cb in self._callbacks:
            cb.epoch_begin(run_context)

    def epoch_end(self, run_context):
        """Execute epoch_end callback"""
        for cb in self._callbacks:
            cb.epoch_end(run_context)

    def step_begin(self, run_context):
        """Execute step_begin callback"""
        for cb in self._callbacks:
            cb.step_begin(run_context)

    def step_end(self, run_context):
        """Execute step_end callback"""
        for cb in self._callbacks:
            cb.step_end(run_context)

    def end(self, run_context):
        """Execute end callback"""
        for cb in self._callbacks:
            cb.end(run_context)

    def append(self, callback: Callback):
        """Add a callback"""
        self._callbacks.append(callback)
        
    def remove(self, callback: Callback):
        """Remove a callback"""
        if callback in self._callbacks:
            self._callbacks.remove(callback)
            
    def __enter__(self):
        """Context manager entry"""
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit"""
        return False 